#!/usr/bin/python
# -*- coding:utf-8 -*-
# @Time  : 6/3/20
# @Author: yanlh
# @usage : 主机配置审计
import sys
import docker_utils

sys.path.append('..')
from common_utils import CBBCommonUtils


def get_plugin_info():
    plugin_info = {
        "name": "Docker_host_001 auditl",
        "plugin_id": "docker_host_001",
        "plugin_type": "Docker_Host_Audit_Safe",
        "info": "Config audit strategies",
        "level": "B",
        "module": "Safety scan",
        "author": "yanlh",
        "keyword": "Safety scan",
        "configable": "false",
    }
    return plugin_info


logger = None
cur_user = None
cur_module = None
cur_task = None


def set_plugin_logger(setter, user, module, task, *args, **kwargs):
    global logger, cur_user, cur_module, cur_task
    logger = setter
    cur_user = user
    cur_module = module
    cur_task = task


rules = {
    'Docker Daemon': '/usr/bin/dockerd',
    '/var/lib/docker': '/var/lib/docker',
    '/etc/docker': '/etc/docker',
    'docker.service': ['docker.service', '-w /usr/lib/systemd/system/docker.service -k docker'],
    'docker.socket': ['docker.socket', '-w /usr/lib/systemd/system/docker.socket -k docker'],
    '/etc/default/docker': '/etc/default/docker',
    '/etc/sysconfig/docker': '/etc/sysconfig/docker',
    '/etc/docker/daemon': '/etc/docker/daemon',
    '/user/bin/containerd': '/usr/bin/containerd',
    '/user/sbin/runc': '/usr/sbin/runc',
}


def single_logger_record(level, des, module='_Scan'):
    if level == 'error':
        logger.debug_error(cur_user, cur_module, cur_task + module, '', des)
    if level == 'info':
        logger.debug_info(cur_user, cur_module, cur_task + module, '', des)
    if level == 'warning':
        logger.debug_warning(cur_user, cur_module, cur_task + module, '', des)


def record_logger(record_list, module):
    des_list = []
    for item in record_list:
        if item['type'] == 'info':
            logger.debug_info(cur_user, cur_module, cur_task + module, '', item['des'])
        if item['type'] == 'error':
            logger.debug_error(cur_user, cur_module, cur_task + module, '', item['des'])
        if item['type'] == 'warning':
            logger.debug_warning(cur_user, cur_module, cur_task + module, '', item['des'])
        des_list.append(item['des'])
    return des_list


def is_auditctl_installed(func_args):
    cmd = 'auditctl'
    result, output = CBBCommonUtils.cbb_run_cmd(func_args['ip'], cmd, username=func_args['username'],
                                                passwd=func_args['passwd'])
    if result:
        des = "No auditctl command or some error happens, error detail {}".format(result)
        single_logger_record('error', des)
        return False, des
    return True, None


def is_audit_rules_exist(func_args, variable, key):
    cmd = 'auditctl -l | grep {}'.format(variable)
    result, output = CBBCommonUtils.cbb_run_cmd(func_args['ip'], cmd, username=func_args['username'],
                                                passwd=func_args['passwd'])
    error = 0
    fail = 0
    if result:
        des = "Not safe. CMD {} process ERROR. Error info: {}".format(variable, result)
        single_logger_record('error', des)
        error += 1
        return des, error, fail
    if output:
        des = 'Safe. {} has configured audit rules.'.format(key)
        single_logger_record('info', des)
    else:
        des = 'Not Safe. {} has not configured any audit rules.'.format(key)
        single_logger_record('warning', des)
        error += 1
        fail += 1
    return des, error, fail


def handle_fragment_file(func_args, variable, key):
    target = '/usr/lib/systemd/system/' + variable
    result, des = docker_utils.validate_fragment_file_exist(func_args, {'path': target})
    if result == 'error':
        return des, 1
    if result == 'fail':
        return des, 0
    des, error, fail = is_audit_rules_exist(func_args, variable, key)
    return des, error, fail


def scan_audit_rule(func_args):
    des_list = []
    error_count = 0
    error_key = []
    for item in rules:
        variable = rules[item]
        if isinstance(variable, list):
            des, error, fail = handle_fragment_file(func_args, variable[0], item)
        else:
            des, error, fail = is_audit_rules_exist(func_args, variable, item)
        des_list.append(des)
        error_count = error_count + error
        if fail:
            error_key.append(item)
    return des_list, error_count, error_key


# 扫描函数
def scan(func_args, flag=0):
    single_logger_record('info', 'Scan Start.')
    installed, des = is_auditctl_installed(func_args)
    if not installed:
        return [des], 1
    des_list, error_count, error_key = scan_audit_rule(func_args)
    des_list.append('Scan Complete.')
    single_logger_record('info', 'Scan Complete.')
    if flag:
        if len(error_key):
            rein_list, error_count = reinforce(func_args, error_key)
            des_list.extend(rein_list)
        single_logger_record('info', '_Reinforce', 'Reinforce Complete.')
    return des_list, error_count


def add_audit_rules(func_args, error_key):
    path = func_args.get('path', None)
    flag = True
    error = 0
    for item in error_key:
        value = rules[item]
        if isinstance(value, list):
            value = rules[item][1]
        else:
            value = '-w {} -k docker'.format(value)
        cmd = "echo '{}' >> {}".format(value, path)
        result, output = CBBCommonUtils.cbb_run_cmd(func_args['ip'], cmd, username=func_args['username'],
                                                    passwd=func_args['passwd'])
        if result:
            flag = False
            des = "Not safe. Add audit rules {} Failed. Error info: {}".format(value, result)
            single_logger_record('warning', '_Reinforce', des)
        else:
            des = 'Safe. Add audit rules {} Success. '.format(value)
            single_logger_record('info', '_Reinforce', des)
    if not flag:
        des = "Not safe. Add audit rules Failed. "
        error = error + 1
    else:
        des = "Safe. Add audit rules Success."
    return des, error


def restart_audit_service(func_args):
    cmd = 'service auditd restart'
    result, output = CBBCommonUtils.cbb_run_cmd(func_args['ip'], cmd, username=func_args['username'],
                                                passwd=func_args['passwd'])
    if result:
        start_cmd = 'service auditd start'
        start_result, start_output = CBBCommonUtils.cbb_run_cmd(func_args['ip'], start_cmd,
                                                                username=func_args['username'],
                                                                passwd=func_args['passwd'])
        if start_result and '/bin/systemctl start auditd.service' not in start_result:
            des = "Not safe. Restart auditd service Failed. Error info: {}".format(result)
            single_logger_record('warning', '_Reinforce', des)
            return des, 1
    des = "Safe. Restart auditd service success."
    return des, 0


def reinforce(func_args, error_key):
    single_logger_record('info', "Reinforce Start.", '_Reinforce')
    result, des = docker_utils.validate_catalog_file_exist(func_args)
    if result != 'success':
        single_logger_record('warning', des, '_Reinforce')
        return [des], 1
    des, error = docker_utils.backup_audit_rules(func_args)
    if error > 0:
        single_logger_record('warning', des, '_Reinforce')
        return [des], 1
    else:
        single_logger_record('info', des, '_Reinforce')
    des_list = []
    des, error = add_audit_rules(func_args, error_key)
    des_list.append(des)
    if error:
        roll_list, roll_error = rollback(func_args)
        des_list.extend(roll_list)
    else:
        restart_des, restart_error = restart_audit_service(func_args)
        des_list.append(restart_des)
        error = error + restart_error
        if restart_error:
            roll_list, roll_error = rollback(func_args)
            des_list.extend(roll_list)
    des_list.append('Reinforce Completed.')
    return des_list, error


# 回滚函数
def rollback(func_args):
    single_logger_record('info', "Rollback Start.", '_Rollback')
    roll_list = []
    record_roll, error_count = docker_utils.rollback_audit_rules(func_args)
    restart_des, restart_error = restart_audit_service(func_args)
    single_logger_record('info', "Rollback Complete.", '_Rollback')
    roll_list.append(record_roll)
    roll_list.append(restart_des)
    if error_count > 0:
        single_logger_record('warning', record_roll, '_Rollback')
    else:
        single_logger_record('info', record_roll, '_Rollback')
    return roll_list, error_count


def check(ip, *args, **kwargs):
    sys_user = kwargs.get("system_user")
    sys_pwd = kwargs.get("system_pwd")
    comm = kwargs.get("command")
    try:
        single_logger_record('info', "docker_host_001 Start.", '_Check')
        func_args = {'ip': ip, 'username': sys_user, 'passwd': sys_pwd, 'comm': comm,
                     'path': '/etc/audit/rules.d/audit.rules'}
        result = docker_utils.check_func(scan, rollback, func_args)
        single_logger_record('info', "docker_host_001 Complete.", '_Check')
        return result
    except Exception as er:
        code = 1
        des = ["ERROR:", str(er)]
        single_logger_record('error', des, '_Check')
        return {"code": code, "count": 0, "des": des}


if __name__ == '__main__':
    """
    本地测试
    """
    check(ip=None, command=2)
